# components/retrievers.py

from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core import VectorStoreIndex
from llama_index.core.schema import NodeWithScore
from typing import List
from ..config import settings


def create_vector_retriever(index: VectorStoreIndex):
    """Creates a vector retriever from the index."""
    print("🛠️ [Retrievers] Creating Vector Retriever...")
    return index.as_retriever(similarity_top_k=settings.VECTOR_SIMILARITY_TOP_K)

def create_bm25_retriever(nodes: List[NodeWithScore]):
    """Creates a BM25 retriever from a list of nodes."""
    print("🛠️ [Retrievers] Creating BM25 Retriever...")
    return BM25Retriever.from_defaults(
        nodes=nodes,
        similarity_top_k=settings.BM25_SIMILARITY_TOP_K
    )

def create_fusion_retriever(vector_retriever, bm25_retriever):
    """Creates a QueryFusionRetriever by combining other retrievers."""
    print("✨ [Retrievers] Creating Fusion Retriever...")
    return QueryFusionRetriever(
        [vector_retriever, bm25_retriever],
        similarity_top_k=settings.FUSION_SIMILARITY_TOP_K,
        num_queries=settings.FUSION_NUM_QUERIES,
        mode=settings.FUSION_MODE,
        use_async=True,
        verbose=True,
    )